import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

print('----实验一-----')
q = np.array([2.321E-19, 2.309E-19, 1.868E-19])
e = 1.6022E-19
n = np.round(q / e)
print(f"算得n_i分别为：{n}")
e_i = q / n
print(f"算得e_i分别为：{e_i}(单位：C）")
e_experient = np.sum(e_i) / 3
print(f"实验所得元电荷电荷量为：{e_experient}(单位：C)\n相对误差为:{abs(e_experient - e) / e * 100}%")
# ————————————————————————实验二——————————————————————————————————————————
print('----实验二-----')
lambda_list = np.array([365, 405, 436, 546, 577]) * 1.0E-9
c = 3.0E+8
v = c / lambda_list
print(f"频率分别为{v / 1.0E14}(单位：10^14Hz)")
U_0 = np.array([1.587, 1.334, 1.058, 0.519, 0.393])
k, b = np.polyfit(v, U_0, 1)
h_experiment = k * e
h = 6.62607015E-34
print(f"k:{k},b:{b},实验测得普朗克常数为{h_experiment}\n相对误差为{abs(h_experiment - h) / h * 100}%")
plt.scatter(v, U_0)
plt.plot(v, k * v + b, color='red')
plt.xlabel("v/10^14Hz")
plt.ylabel("U_O/V")
plt.title("Experiment Two")
plt.savefig('C:/Users/admin/Desktop/实验二.png', dpi=500, bbox_inches='tight')
plt.show()
# ————————————————————————实验三——————————————————————————————————————————
filename = "E:/data.csv"
df = pd.read_csv(filename)
U_AK_1 = np.array(df.iloc[0, 0:3].values)
U_AK_2 = np.array(df.iloc[0, 2:8].values)
U_AK_3 = np.array(df.iloc[0, 7:28].values)
color_list = ['b', 'm', 'g', 'y', 'r']
figure_data_list = []


def fun(i):
    I_1 = np.array(df.iloc[i, 0:3].values)
    coefficients = np.polyfit(U_AK_1, I_1, 3)
    polynomial = np.poly1d(coefficients)
    x_values = np.linspace(min(U_AK_1), max(U_AK_1), 100)
    y_values = polynomial(x_values)
    plt.plot(U_AK_1, I_1, 'o')
    plt.plot(x_values, y_values, '-', label=f'lambda={int(lambda_list[i - 1] * 1.0E+9)}mm')
    I_2 = np.array(df.iloc[i, 2:8].values)
    coefficients = np.polyfit(U_AK_2, I_2, 3)
    polynomial = np.poly1d(coefficients)
    x_values = np.linspace(min(U_AK_2), max(U_AK_2), 100)
    y_values = polynomial(x_values)
    plt.plot(U_AK_2, I_2, 'o')
    plt.plot(x_values, y_values, '-', label=f'lambda={int(lambda_list[i - 1] * 1.0E+9)}mm')
    I_3 = np.array(df.iloc[i, 7:28].values)
    coefficients = np.polyfit(U_AK_3, I_3, 3)
    polynomial = np.poly1d(coefficients)
    x_values = np.linspace(min(U_AK_3), max(U_AK_3), 100)
    y_values = polynomial(x_values)
    plt.plot(U_AK_3, I_3, 'o', color=color_list[i - 1])
    p, = plt.plot(x_values, y_values, '-', color=color_list[i - 1],
                    label=f'lambda={int(lambda_list[i - 1] * 1.0E+9)}mm')
    figure_data_list.append(p)


# 三次多项式拟合。观察数据可以得出，以U_AK=0，5为界，函数的趋势差异较为明显，故进行分段拟合
for i in range(1, 6):
    fun(i)

plt.xlabel("U_AK/V")
plt.ylabel("I/A")
plt.title("Experiment Three\nI*E-10 for lambda = 365,405,436;I*E-11 for lambda = 546,577")
plt.legend([figure_data_list[0], figure_data_list[1], figure_data_list[2], figure_data_list[3], figure_data_list[4]], [f"lambda={int(lambda_list[0] * 1.0E+9)}mm", f"lambda={int(lambda_list[1] * 1.0E+9)}mm", f"lambda={int(lambda_list[2] * 1.0E+9)}mm", f"lambda={int(lambda_list[3] * 1.0E+9)}mm", f"lambda={int(lambda_list[4] * 1.0E+9)}mm"])
plt.savefig('C:/Users/admin/Desktop/实验三.png', dpi=500, bbox_inches='tight')
plt.show()
